from ultralytics import YOLO
import time
import cv2
import os
import numpy as np



def image_resize(image, resize_rate):
    src_w = image.shape[0]
    src_h = image.shape[1]

    # image = cv.resize(image, None, fx=resize_rate, fy=resize_rate, interpolation=cv.INTER_LINEAR)
    new_image = cv2.resize(image, None, fx=resize_rate, fy=resize_rate, interpolation=cv2.INTER_NEAREST)
    return new_image


def list_dir(path, file_list, file_type='txt', filter_str=None):
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if os.path.isdir(file_path):
            list_dir(file_path, file_list, file_type, filter_str=filter_str)
        else:
            if filter_str is None:
                if file_path[-len(file_type):].lower() == file_type.lower():
                    file_list.append(file_path)
            else:
                if filter_str not in file_path:
                    if file_path[-len(file_type):].lower() == file_type.lower():
                        file_list.append(file_path)



def loc_(img, model):
    # result = model.predict(source='/home/yimu/Backup/projects_zy/datasets/hangtian_kegong/dianrong_noclass/images/test',save=False)
    #
    count = 0
    result = model(img)
    target_list = []
    for r in result:

        # print(r.boxes)
        for re in r.boxes:
            xyxy = re.cpu().numpy().xyxy.astype(int)
            conf = re.cpu().numpy().conf[0]
            # print(xyxy[0][1],xyxy[0][3], xyxy[0][0],xyxy[0][2])
            crop_img = img[xyxy[0][1]:xyxy[0][3], xyxy[0][0]:xyxy[0][2]]
            # print(count, crop_img.shape[0] * crop_img.shape[1])
            px = int((xyxy[0][0] + xyxy[0][2]) / 2)
            py = int((xyxy[0][1] + xyxy[0][3]) / 2)

            # cv2.circle(img, (px, py), radius=5, color=(0, 0, 255), thickness=-1)
            target = [conf, px, py]
            target_list.append(target)
            # print(conf)
            if crop_img.shape[0] < 414 and crop_img.shape[1] < 414:
                # res = Watershed(crop_img)
                # cv2.imshow('box', res)
                # cv2.imshow('box2', crop_img)
                # cv2.rectangle(img, (xyxy[0][0], xyxy[0][1]), (xyxy[0][2], xyxy[0][3]), (0, 0, 255), 5)
                # cv2.imshow('box', image_resize(img, 0.5))
                # cv2.waitKey()
                # print(count)
                # cv2.imwrite(
                #     '/home/yimu/Backup/projects_zy/datasets/hangtian_kegong/dianrong_noclass/images/test_save/' + str(
                #         count) + '.jpg', crop_img)
                count += 1
    # cv2.imshow('box', img)
    # cv2.waitKey()
    return target_list

if __name__ == '__main__':

    model = YOLO('best.pt')  # load a pretrained model (recommended for training)
    model.info()
    img = cv2.imread('test_pic/Image_20231228135936019.bmp')
    a = loc_(img, model)
    print(a)
    # file_path = 'test_pic/'
    # file_list = []
    # list_dir(file_path, file_list, 'bmp')
    # count = 0
    # for img_path in file_list:
    #
    #     #
    #     result = model(img_path)
    #     img = cv2.imread(img_path)
    #     target_list = []
    #     for r in result:
    #
    #         # print(r.boxes)
    #         for re in r.boxes:
    #             xyxy = re.cpu().numpy().xyxy.astype(int)
    #             conf = re.cpu().numpy().conf[0]
    #             crop_img = img[xyxy[0][1]:xyxy[0][3], xyxy[0][0]:xyxy[0][2]]
    #             # print(count, crop_img.shape[0] * crop_img.shape[1])
    #             px = int((xyxy[0][0] + xyxy[0][2]) / 2)
    #             py = int((xyxy[0][1] + xyxy[0][3]) / 2)
    #
    #             cv2.circle(img, (px, py), radius=5, color=(0, 0, 255), thickness=-1)
    #             target = [conf, px, py]
    #             target_list.append(target)
    #             print(conf)
    #             if crop_img.shape[0] < 414 and crop_img.shape[1] < 414:
    #                 # res = Watershed(crop_img)
    #                 # cv2.imshow('box', res)
    #                 # cv2.imshow('box2', crop_img)
    #                 cv2.rectangle(img, (xyxy[0][0], xyxy[0][1]), (xyxy[0][2], xyxy[0][3]), (0, 0, 255), 5)
    #                 cv2.imshow('box', image_resize(img, 0.5))
    #                 cv2.waitKey()
    #                 print(count)
    #                 # cv2.imwrite(
    #                 #     '/home/yimu/Backup/projects_zy/datasets/hangtian_kegong/dianrong_noclass/images/test_save/' + str(
    #                 #         count) + '.jpg', crop_img)
    #                 count += 1
    #     cv2.imshow('box', img)
    #     cv2.waitKey()
    #     # print(result[0])
